#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>

#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libswscale/swscale.h>

#include "encode.h"

typedef struct {
    int size;
    uint8_t* buffer;
} pngout_t;

static void SaveFrame(AVFrame *pFrame, int width, int height, FILE* pFile) {
    int y;

    // Write header
    fprintf(pFile, "P6\n%d %d\n255\n", width, height);

    // Write pixel data
    for (y = 0; y < height; y++)
        fwrite(pFrame->data[0] + y * pFrame->linesize[0], 1, width * 3, pFile);

}

static void SaveExternal(pngout_t* image, char* filename, char* output_format) {
    FILE *pFile;
    char szFilename[200];

    // Open file
    sprintf(szFilename, "%s.%s", filename, output_format);
    pFile = fopen(szFilename, "wb");
    if (pFile == NULL)
        return;

    // Write pixel data
    fwrite(image->buffer, 1, image->size, pFile);
    // Close file
    fclose(pFile);
}

pngout_t* frame_to_external(AVFrame* pFrame, int width, int height, char* format) {
    //frame_to_png(AVFrame *pFrame, int width, int height, char* format) {
    AVCodecContext *encoder_codec_ctx = NULL;
    AVCodec *encoder_codec = NULL;

    uint8_t *encoder_output_buffer = NULL;
    size_t encoder_output_buffer_size;
    int err;
    pngout_t* ret = malloc(sizeof (pngout_t));

    fflush(stdout);

    encoder_codec = avcodec_find_encoder_by_name(format);
    if (encoder_codec == NULL) {
        printf("Couldn't find a PNG encoder\n");
        ret->size = -1;
        return ret;
    }

    //TODO: no efficient as we allocate a maximum of bytes
    encoder_output_buffer_size = 100000000 * sizeof (int);
    encoder_output_buffer = av_malloc(encoder_output_buffer_size);
    if (encoder_output_buffer == NULL) {
        printf("Failed to allocate the encoder output buffer\n");
        ret->size = -1;
        return ret;
    }

    encoder_codec_ctx = avcodec_alloc_context();
    if (encoder_codec_ctx == NULL) {
        printf("Failed to allocate the encoder codec context\n");
        av_free(encoder_output_buffer);
        ret->size = -1;
        return ret;
    }

    encoder_codec_ctx->width = width;
    encoder_codec_ctx->height = height;
    encoder_codec_ctx->pix_fmt = PIX_FMT_RGB24;

    if (avcodec_open(encoder_codec_ctx, encoder_codec) < 0) {
        printf("Failed to open the encoder\n");
        av_free(encoder_codec_ctx);
        av_free(encoder_output_buffer);
        ret->size = -1;
        return ret;
    }

    err = avcodec_encode_video(encoder_codec_ctx,
            encoder_output_buffer,
            encoder_output_buffer_size, pFrame);

    avcodec_close(encoder_codec_ctx);
    av_free(encoder_codec_ctx);

    ret->size = err;
    ret->buffer = encoder_output_buffer;


    fflush(stdout);
    return ret;
}

/*
int encode_to_j2k(AVFrame* pFrameRGB, int width, int height, char* filename, char* output_format) {
    if (strcmp("raw", output_format) == 0) {
        SaveFrame(pFrameRGB, width, height, filename);
    } else {
        pngout_t* ret = malloc(sizeof (pngout_t));
        ret->buffer = NULL;
        ret = frame_to_external(pFrameRGB, width, height, output_format);
        SaveExternal(ret, filename, output_format);
        if (ret->buffer != NULL) {
            free(ret->buffer);
        }
        free(ret);
    }
    return 0;
}
*/

int saveFrameToDisk(AVFrame* pFrameRGB, int width, int height, char* filename, char* output_format) {
    if (strcmp("raw", output_format) == 0) {
        FILE *pFile;
        char szFilename[55];
        // Open file
        sprintf(szFilename, "%s.ppm", filename);

        pFile = fopen(szFilename, "wb");
        if (pFile == NULL)
            return;
        SaveFrame(pFrameRGB, width, height, pFile);
        fclose(pFile);
        // Close file

    } else {
        pngout_t* ret = malloc(sizeof (pngout_t));
        ret->buffer = NULL;
        ret = frame_to_external(pFrameRGB, width, height, output_format);
        SaveExternal(ret, filename, output_format);
        if (ret->buffer != NULL) {
            free(ret->buffer);
        }
        free(ret);
    }
    return 0;
}

int encodeFrameInJ2K(AVFrame* pFrame, int width, int height, char* filename, char* output_format) {
  /*char szFilename[55];
    sprintf(szFilename, "%s.ppm", filename);


    
    saveFrameToDisk(pFrame, width, height, filename, "raw");
    //    encode_from_file(char*file, char* outfile, char* comment, int cinema_format, char* format)
    encode_from_file(szFilename, szFilename2, "", 0, "raw");
    */
    


    FILE *pFile;
    char szFilename2[55];
    sprintf(szFilename2, "%s.j2k", filename);
    void *data = malloc(10000000 * sizeof (int));

    pFile = fmemopen(data, 10000000 * sizeof (int), "wb");
    if (data == NULL)
        return;
    SaveFrame(pFrame, width, height, pFile);
    fclose(pFile);

/*
    encode_from_image(data, 10000000 * sizeof (int), szFilename2, "", 0);
*/

    // Close file
    free(data);


    /*
        int y;
        int x;
        char *data = malloc(width * height * 3 * sizeof(char));
        for (y = 0; y < height; y++) {
            for (x = 0; x < width; x++) {
     *data++ = (pFrame->data[0] + y * pFrame->linesize[0]) + x;
            }
        }

       
        return 0;
     */
}

int decode_encode_local(char * file, char* output_dir, int num_threads) {
    _decode(file, output_dir, "ppm", &encodeFrameInJ2K);
}

int decode_only(char * file, char* output_dir, char*output_format) {
    _decode(file, output_dir, output_format, &saveFrameToDisk);
}

int _decode(char * file, char* output_dir, char* output_format, int(*frame_decoded)(AVFrame*, int, int, char*, char*)) {
    AVFormatContext *pFormatCtx;
    int i, videoStream;
    AVCodecContext *pCodecCtx;
    AVCodec *pCodec;
    AVFrame *pFrame;
    AVFrame *pFrameRGB;
    AVPacket packet;
    int frameFinished;
    int numBytes;
    uint8_t *buffer;

    int out_width = 2048;
    int out_height = 1080;

    // Register all formats and codecs
    av_register_all();

    // Open video file
    if (av_open_input_file(&pFormatCtx, file, NULL, 0, NULL) != 0)
        return -1; // Couldn't open file

    // Retrieve stream information
    if (av_find_stream_info(pFormatCtx) < 0)
        return -1; // Couldn't find stream information

    // Dump information about file onto standard error
    dump_format(pFormatCtx, 0, file, false);

    // Find the first video stream
    videoStream = -1;
    for (i = 0; i < pFormatCtx->nb_streams; i++)
        if (pFormatCtx->streams[i]->codec->codec_type == CODEC_TYPE_VIDEO) {
            videoStream = i;
            break;
        }
    if (videoStream == -1)
        return -1; // Didn't find a video stream

    // Get a pointer to the codec context for the video stream
    pCodecCtx = pFormatCtx->streams[videoStream]->codec;

    // Find the decoder for the video stream
    pCodec = avcodec_find_decoder(pCodecCtx->codec_id);
    if (pCodec == NULL)
        return -1; // Codec not found

    // Open codec
    if (avcodec_open(pCodecCtx, pCodec) < 0)
        return -1; // Could not open codec

    // Hack to correct wrong frame rates that seem to be generated by some codecs
    if (pCodecCtx->time_base.num > 1000 && pCodecCtx->time_base.den == 1)
        pCodecCtx->time_base.den = 1000;

    // Allocate video frame
    pFrame = avcodec_alloc_frame();

    // Allocate an AVFrame structure
    pFrameRGB = avcodec_alloc_frame();
    if (pFrameRGB == NULL)
        return -1;

    // Determine required buffer size and allocate buffer
    numBytes = avpicture_get_size(PIX_FMT_RGB24, out_width,
            out_height);

    buffer = malloc(numBytes);

    // Assign appropriate parts of buffer to image planes in pFrameRGB
    avpicture_fill((AVPicture *) pFrameRGB, buffer, PIX_FMT_RGB24,
            out_width, out_height);

    // Read frames and save first five frames to disk
    i = 0;
    static struct SwsContext * img_convert_ctx;
    while (av_read_frame(pFormatCtx, &packet) >= 0) {
        // Is this a packet from the video stream?
        if (packet.stream_index == videoStream) {
            // Decode video frame
            avcodec_decode_video2(pCodecCtx, pFrame, &frameFinished, &packet);

            // Did we get a video frame?
            if (frameFinished) {
                // Convert the image into YUV format that SDL uses
                if (img_convert_ctx == NULL) {
                    img_convert_ctx = sws_getContext(pCodecCtx->width,
                            pCodecCtx->height,
                            pCodecCtx->pix_fmt,
                            out_width,
                            out_height,
                            PIX_FMT_RGB24,
                            SWS_BICUBIC, (int) NULL, NULL, NULL);
                    //This way should normally work
                    /*img_convert_ctx = sws_alloc_context();

                    av_set_int(img_convert_ctx, "srcw", pCodecCtx->width);
                    av_set_int(img_convert_ctx, "srch", pCodecCtx->height);
                    av_set_int(img_convert_ctx, "dstw", pCodecCtx->width);
                    av_set_int(img_convert_ctx, "dsth", pCodecCtx->height);
                    av_set_int(img_convert_ctx, "src_format", pCodecCtx->pix_fmt);
                    av_set_int(img_convert_ctx, "dst_format", PIX_FMT_RGB24);
                    av_set_int(img_convert_ctx, "sws_flags", SWS_BICUBIC);

                    av_set_int(img_convert_ctx, "src_range", NULL);
                    av_set_int(img_convert_ctx, "dst_range", NULL);
                     */

                    if (sws_init_context(img_convert_ctx, NULL, NULL) < 0) {
                        fprintf(stderr, "Cannot initialize the conversion context!\n");
                        exit(1);
                    }
                }
                sws_scale(img_convert_ctx, (const uint8_t * const*) pFrame->data, pFrame->linesize, 0, pCodecCtx->height, pFrameRGB->data, pFrameRGB->linesize);

                char filename[55];
                sprintf(filename, "%sframe%09d", output_dir, i);

                if (((*frame_decoded)(pFrameRGB, out_width, out_height, filename, output_format)) < 0) {
                    printf("Error during encoding frame : %d\n", i);
                }

                /*	    if (strcmp("raw", output_format) == 0) {
                  SaveFrame(pFrameRGB, out_width, out_height, output_dir, i);
                } else {
                  ret = frame_to_png(pFrameRGB, pCodecCtx->width, pCodecCtx->height, output_format);
                  SavePng(ret, output_dir, output_format, i);
                  if (ret->buffer != NULL) {
                    free(ret->buffer);
                    }*/
            }
            i++;
        }
        // Free the packet that was allocated by av_read_frame
        av_free_packet(&packet);
    }

    free(img_convert_ctx);

    // Free the RGB image
    free(buffer);
    av_free(pFrameRGB);

    // Free the YUV frame
    av_free(pFrame);

    // Close the codec
    avcodec_close(pCodecCtx);

    // Close the video file
    av_close_input_file(pFormatCtx);

    return i;
}
